{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "62549efe",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Examples\n",
    "\n",
    "`Kork` allows providing examples in the prompt. \n",
    "\n",
    "1. Examples help the LLM understand the syntax of the language\n",
    "2. Examples can be dynamically generated based on the user query, to surface the most relevant programs.\n",
    "\n",
    "The format of examples is a sequence of 2-tuples of the form: (`str`, `Program`).\n",
    "\n",
    "Let's see how to specify them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "79fd8e41",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "nbsphinx": "hidden",
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": [
     "remove_cell"
    ]
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "\n",
    "sys.path.insert(0, \"../\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27975824-3db3-48ad-bbdd-6784827f7a99",
   "metadata": {},
   "source": [
    "## The Simplest Way"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5c5b3a8d-c464-499c-a5d1-275ab2ceb226",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kork.parser import parse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1a6dfced-5e2c-4c34-8dcd-a5d207bcd92d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "examples_as_strings = [\n",
    "    (\n",
    "        \"declare a variable called `y` and assign to it the value 8\",\n",
    "        \"var y = 8\",\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "23b683ed-f52a-4fce-b554-9fbbbaa4af4c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "examples = [(query, parse(code)) for query, code in examples_as_strings]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "03e16c63-f804-4076-a70a-2d324dc43e8a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('declare a variable called `y` and assign to it the value 8',\n",
       "  Program(stmts=(VarDecl(name='y', value=Literal(value=8)),)))]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fb92038-08ad-4cf6-91db-06fb25c281bd",
   "metadata": {
    "tags": []
   },
   "source": [
    "## The Hardest Way\n",
    "\n",
    "Let's import AST elements and write an example query and a corresponding short program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b5775536-ea67-4ed6-97e7-4d6f69f8d83a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kork.ast import Program, VarDecl, Literal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b6fec8a2-e4ce-48a8-b490-1e8be83886e4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "examples = [\n",
    "    (\n",
    "        \"declare a variable called `y` and assign to it the value 8\",\n",
    "        Program(stmts=[VarDecl(\"y\", Literal(value=8))]),\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b8b7689-2d41-4693-a389-7ccf5916aeb3",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Formatted Example\n",
    "\n",
    "The chain formats the example using an `AstPrinter` to format the examples as a sequence of the 2-tuple of (`str`, `str`).\n",
    "\n",
    "Here's what happens:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a1e2b3c0-e049-49c2-af20-d5c52e34db9c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kork.examples import format_examples\n",
    "from kork import AstPrinter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "83eea564-0102-4d89-b627-7781d4d76a53",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('declare a variable called `y` and assign to it the value 8', '```SmirkingCat\\nvar y = 8\\n```')]\n"
     ]
    }
   ],
   "source": [
    "print(format_examples(\"SmirkingCat\", examples, AstPrinter()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2e8c20f7-b419-4955-b5e3-c8046ec29357",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('Text: \"\"\"\\ndeclare a variable called `y` and assign to it the value 8\\n\"\"\"', '```SmirkingCat\\nvar y = 8\\n```')]\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    format_examples(\n",
    "        \"SmirkingCat\", examples, AstPrinter(), input_formatter=\"text_prefix\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0db700c-f182-4c24-a060-2e19f4b3102d",
   "metadata": {},
   "source": [
    "## The Hardest Way Made Easier\n",
    "\n",
    "There are utility functions that make it easier to generate programs that invoke foreign functions -- which is\n",
    "expected to be a common pattern."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0fe1ff0d-f703-47a6-8398-97c5a853710f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kork import c_, r_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fb84132e-a6ac-4b47-a891-2667420b24c3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m \u001b[0mc_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mCallable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mkork\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mast\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mFunctionCall\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m Create a kork function call.\n",
       "\u001b[0;31mFile:\u001b[0m      ~/src/kork/kork/examples.py\n",
       "\u001b[0;31mType:\u001b[0m      function"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "?c_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b5f678c5-4715-4e96-950a-9b82ed0485d4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m \u001b[0mr_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexpr\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mkork\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mast\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mExpr\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mkork\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mast\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mProgram\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m Assign last program expression to a result variable.\n",
       "\u001b[0;31mFile:\u001b[0m      ~/src/kork/kork/examples.py\n",
       "\u001b[0;31mType:\u001b[0m      function"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "?r_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8aeef131-2c87-473c-9780-5d6e391993aa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import operator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "054f4260-7c49-4af9-8720-25c990379ee9",
   "metadata": {
    "tags": []
   },
   "source": [
    "Here's a small program that invokes the foreign function `log10` with arg `2`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3b411f83-4413-4ace-a786-5c9a13bca343",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionCall(name='log10', args=[Literal(value=2)])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c_(math.log10, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86e9bee5-558d-4344-a315-29770d9b1b3e",
   "metadata": {
    "tags": []
   },
   "source": [
    "A program that invokes `log10` with arg of 2 and stores the result in a `result` variable in the global namespace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "07480868-5d28-48e1-8b58-1b72c0c65ee0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Program(stmts=[VarDecl(name='result', value=FunctionCall(name='log10', args=[Literal(value=2)]))])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r_(c_(math.log10, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d4113fd1-f17d-4177-a18d-e571088d329e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "program = r_(c_(operator.add, 2, c_(math.log10, 3)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9facd1c6-b203-49d1-840e-12a3efaa5b11",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "var result = add(2, log10(3))\n"
     ]
    }
   ],
   "source": [
    "print(AstPrinter().visit(program))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "77fcb89f-a8bb-4ff6-a63b-574f39087642",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "examples = [(\"take the log base 10 of 3 and add 2 to the result.\", program)]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
